#include "utils.h"

template<int bt_sz, int bn>
__global__ void double_buffer_kernel(
  const float * __restrict__ src, float * __restrict__ dst, int sz, const float val)
{
  // float2 rtmp;
  // int tx = threadIdx.x, bx = blockIdx.x;
  // int b_start = bx*bt_sz*bn;
  // for(int id = 0; id<(bn+1); ++id){
  //   if(id < bn){
  //     int buf_id = (id+1)%2;
  //     int bn_id = id;
  //     int idx = b_start + bn_id*bt_sz + tx;
  //     if(idx < sz){
  //       ((float*)(&rtmp))[buf_id] = src[idx];
  //     }
  //   }
  //   if(id > 0){
  //     int buf_id = id%2;
  //     int bn_id = id-1;
  //     int idx = b_start + bn_id*bt_sz + tx;
  //     if(idx < sz){
  //       int sum = ((float*)(&rtmp))[buf_id];
  //       for(int i = 0; i < 512; ++i){
  //         sum += val;
  //       }
  //       dst[idx] = sum;
  //     }
  //   }
  // }

  int tx = threadIdx.x, bx = blockIdx.x;
  int b_start = bx*bt_sz*bn;
  for(int bn_id = 0; bn_id < bn; ++bn_id){
    int idx = b_start + bn_id*bt_sz + tx;
    if(idx < sz){
      int sum = src[idx];
      for(int i = 0; i < 512; ++i){
        sum += val;
      }
      dst[idx] = sum;
    }
  }
}
void test_double_buffer(cudaStream_t cuStream, float *src, float *dst, int sz, const float val = 1.0f){
  const int bn = 16;
  const int bt_sz = 128;
  dim3 g_sz(DIV_UP(sz, bt_sz*bn)), b_sz(bt_sz);
  double_buffer_kernel<bt_sz, bn><<<g_sz, b_sz, 0, cuStream>>>(src, dst, sz, val);
}


__global__ void shared_broadcast_kernel(const float * __restrict__ src, float * __restrict__ dst)
{
  int tx = threadIdx.x;
  __shared__ float stmp[512];
  stmp[tx] = src[tx];

  // float2 val = ((float2*)stmp)[tx%16]; //broadcast
  // ((float2*)dst)[tx] = val;
  float4 val = ((float4*)(stmp+(tx/4)*32))[tx%4]; //8 way bank conflict
  ((float4*)dst)[tx] = val;
  // float4 val = ((float4*)stmp)[tx%4]; //broadcast
  // ((float4*)dst)[tx] = val;
  // float4 val = ((float4*)(stmp+(tx/2)*32))[tx%2]; //16 way bank conflict
  // ((float4*)dst)[tx] = val;
}

void test_shared_broadcast(cudaStream_t cuStream, float *src, float *dst, int sz, const float val = 1.0f){
  const int bt_sz = 32;
  dim3 g_sz(DIV_UP(1, bt_sz)), b_sz(bt_sz);
  shared_broadcast_kernel<<<g_sz, b_sz, 0, cuStream>>>(src, dst);
}

__global__ void threadfence_kernel(const float * __restrict__ src, float * __restrict__ dst)
{
  __shared__ int stmp[128];
  int tx = threadIdx.x;

  int warp_id = tx/32, tid = tx%32;
  
  int dst_warp = 1;
  int val = 1001;

  stmp[warp_id*32+tid] = 0;
  __syncthreads();

  if(warp_id == 0){
    stmp[dst_warp*32+tid] = val;
  }
  printf("(%d %d) \n", warp_id, stmp[warp_id*32+tid]);
}

void test_threadfence(cudaStream_t cuStream, float *src, float *dst, int sz, const float val = 1.0f){
  const int bt_sz = 128;
  dim3 g_sz(DIV_UP(1, bt_sz)), b_sz(bt_sz);
  threadfence_kernel<<<g_sz, b_sz, 0, cuStream>>>(src, dst);
}

int main(int argc, char *argv[]) {
    cudaSetDevice(0);
    
    const int loop = 1;
    cudaStream_t cuStream = NULL;
    CUDA_CHECK(cudaStreamCreateWithFlags(&cuStream, cudaStreamNonBlocking));

    float *d_src = nullptr;
    float *d_dst = nullptr;
    int sz = 10000000;

    CUDA_CHECK(cudaMalloc(reinterpret_cast<void **>(&d_src), sizeof(float)*sz));
    CUDA_CHECK(cudaMalloc(reinterpret_cast<void **>(&d_dst), sizeof(float)*sz));

    int start = now();
    for(int i = 0; i < loop; ++i){
        // test_double_buffer(cuStream, d_src, d_dst, sz, 1.0f);
        // test_shared_broadcast(cuStream, d_src, d_dst, sz, 1.0f);
        test_threadfence(cuStream, d_src, d_dst, sz, 1.0f);
    }
    cudaStreamSynchronize(cuStream);
    int end = now();
    
    float avg_latency =float(end-start)/float(loop);
    printf("loop %d avg_latency %0.4f ms  \n", loop, avg_latency);
    // 0.89 ms

    CUDA_CHECK(cudaFree(d_src));
    CUDA_CHECK(cudaFree(d_dst));
    CUDA_CHECK(cudaStreamDestroy(cuStream));
    return 0;
}